using System;
using System.Configuration.Install;
using System.Diagnostics;
using System.ComponentModel;
using System.Collections.Generic;
using System.Text;
using Microsoft.Exchange.Data.Transport.Smtp;
using Microsoft.Win32;

namespace HeaderFilterAgent
{
	[RunInstaller(true)]
	public class CustomEventLogInstaller : Installer
	{
		private EventLogInstaller eventLogInstaller;
		public CustomEventLogInstaller()
		{
			// Create an instance of 'EventLogInstaller'.
			eventLogInstaller = new EventLogInstaller();
			// Set the 'Source' of the event log, to be created.
			eventLogInstaller.Source = "HeaderFilterAgent";
			// Set the 'Event Log' that the source is created in.
			eventLogInstaller.Log = "Application";
			// Add myEventLogInstaller to 'InstallerCollection'.
			Installers.Add(eventLogInstaller);
		}
	}

    public sealed class HeaderFilterAgentFactory : SmtpReceiveAgentFactory
    {
        public override SmtpReceiveAgent CreateAgent(Microsoft.Exchange.Data.Transport.SmtpServer server)
        {
            return new HeaderFilterAgent();
        }
    }

    public class HeaderFilterAgent : SmtpReceiveAgent
    {
		bool logging = false;
		private RegistryKey headerFilterKey;
		private List<string> whitelist = new List<string>();
		private StringComparer stringComparer = StringComparer.OrdinalIgnoreCase;

        public HeaderFilterAgent()
        {
            this.headerFilterKey = Registry.LocalMachine.OpenSubKey("Software\\HeaderFilterAgent");
			if (this.headerFilterKey == null)
				return;

			if (headerFilterKey.GetValue("Logging") != null)
				logging = ((int)headerFilterKey.GetValue("Logging") == 0 ? false : true);

			if (logging)
				EventLog.WriteEntry("HeaderFilterAgent", "Loading Whitelist.", EventLogEntryType.Information, 1);

			// Read the whitelist into a temporary array
			string[] tempList = (string[])headerFilterKey.GetValue("Whitelist");
			if (tempList == null)
				return;

			if (logging)
			{
				StringBuilder logString = new StringBuilder();
				foreach (string header in tempList)
					logString.Append(header + "\n");

				EventLog.WriteEntry("HeaderFilterAgent", "Whitelist loaded: \n" + logString.ToString(), EventLogEntryType.Information, 2);
			}
			// Convert all the listed headers to uppercase and add them to the collection.
			for (int x = 0; x < tempList.Length; x++)
				whitelist.Add(tempList[x].ToUpper());

			this.OnEndOfHeaders += new EndOfHeadersEventHandler(HeaderFilterAgent_OnEndOfHeaders);

			if (logging)
				EventLog.WriteEntry("HeaderFilterAgent", "HeaderFilterAgent finished loading.", EventLogEntryType.Information, 3);
        }

        void HeaderFilterAgent_OnEndOfHeaders(ReceiveMessageEventSource source, EndOfHeadersEventArgs e)
        {
			try
			{
				if (logging)
				{
					StringBuilder logString = new StringBuilder();
					foreach (Microsoft.Exchange.Data.Mime.Header header in e.Headers)
						logString.Append(header.Name + ":" + header.Value + "\n");

					EventLog.WriteEntry("HeaderFilterAgent", "OnEndOfHeaders was called for message with headers: \n" + logString.ToString(), EventLogEntryType.Information, 4);
				}

				if (e.SmtpSession.AuthenticationSource == AuthenticationSource.Anonymous)
				{
					if (logging)
						EventLog.WriteEntry("HeaderFilterAgent", "AuthenticationSource is Anonymous. Headers will be filtered.", EventLogEntryType.Information, 5);

					foreach (Microsoft.Exchange.Data.Mime.Header header in e.Headers)
					{
						if (header.Name.StartsWith("X-", StringComparison.OrdinalIgnoreCase))
						{
							if (logging)
								EventLog.WriteEntry("HeaderFilterAgent", "Checking header: " + header.Name, EventLogEntryType.Information, 6);

							if (!whitelist.Contains(header.Name.ToUpper()))
							{
								if (logging)
									EventLog.WriteEntry("HeaderFilterAgent", "Removing header: " + header.Name, EventLogEntryType.Information, 7);

								header.RemoveFromParent();
							}
						}
					}
				}
			}
			catch (Exception exc)
			{
				EventLog.WriteEntry("HeaderFilterAgent", "Exception: " + exc.Message + "\n" + exc.StackTrace, EventLogEntryType.Error, 8);
			}
			finally
			{
				if (logging)
					EventLog.WriteEntry("HeaderFilterAgent", "Finished processing message.", EventLogEntryType.Information, 9);
			}
        }
    }
}
